import torch
import numpy as np
import deepxde as dde

class FrequencyPredictor():
    def __init__(self, PDECase, plot_v_error=False):
        self.name = "FrequencyPredictor"
        self.PDECase = PDECase
        self.plot_v_error = plot_v_error
        self.v = torch.tensor(PDECase.v)
        #self.v = dde.Variable(PDECase.v)
        net = self.build_net()
        self.model = dde.Model(PDECase.data, net)
        
    def build_net(self):
        layer_size = self.PDECase.FP_size
        activation = self.activation
        initializer = "Glorot normal"
        net = dde.nn.FNN(layer_size, activation, initializer)
        #for name, param in net.named_parameters():
        #    param.requires_grad = False
        return net
    
    def activation(self, x):
        return torch.cos(torch.abs(self.v) * x)
    
    def eval(self):
        X, y = self.PDECase.gen_testdata()
        y_pred = self.model.predict(X)
        error = dde.metrics.l2_relative_error(y, y_pred)
        return error
    
    def get_v_with_min_error(self, start=1.0, end=100.0, step=0.1):
        error = []
        min_error = 1e10
        min_v = 0
        for i, v in enumerate(np.arange(start, end, step)):
            self.v = torch.tensor(v)
            error.append(self.eval())
            if error[-1] < min_error:
                min_error = error[-1]
                min_v = v
        if self.plot_v_error:
            import matplotlib.pyplot as plt
            plt.plot(np.arange(start, end, step), error)
            plt.show()
            print("min_v: ", min_v)
        return torch.tensor(min_v)
    
    def train(self):
        #self.model.compile("adam", lr=0.0001, external_trainable_variables=[self.v])
        #variable = dde.callbacks.VariableValue([self.v], period=600, filename="variables.dat")
        #losshistory, train_state = self.model.train(iterations=self.PDECase.max_iter, callbacks=[variable])
        self.model.compile("adam", lr=0.0001)
        losshistory, train_state = self.model.train(iterations=self.PDECase.max_iter)
        self.v = self.get_v_with_min_error(start=self.v.cpu().numpy()*10, end=1000.0, step=0.1)



class cFPCT_DNN():
    def __init__(self, PDECase):
        self.name = "cFPCT_DNN"
        self.PDECase = PDECase
        self.FP = FrequencyPredictor(PDECase, plot_v_error=True)
        self.model = None
        self.v = 0.0

    def build_net(self):
        layer_size = self.PDECase.DNN_size
        activation = self.activation(self.v)
        initializer = "Glorot normal"
        net = dde.nn.FNN(layer_size, activation, initializer)
        return net
    
    def activation(self, v):
        def func(x):
            return torch.cos(v * x)
        return func
    
    def train(self):
        self.FP.train()
        net_num = max(1, len(self.PDECase.DNN_size) - 2)
        self.v = torch.pow(self.FP.v, 1/net_num).mul(self.PDECase.DNN_size[0]) .detach()
        self.model = dde.Model(self.PDECase.data, self.build_net())
        self.model.compile("adam", lr=0.0001)
        losshistory, train_state = self.model.train(iterations=self.PDECase.max_iter)